from typing import Optional, List, Dict, Any
from datetime import datetime
from bson import ObjectId

class UserModel:
    """用户模型"""
    def __init__(self, username: str, password: str, email: str, phone: Optional[str] = None,
                 is_admin: bool = False, avatar: Optional[str] = None, 
                 created_at: Optional[datetime] = None, _id: Optional[ObjectId] = None):
        self._id = _id or ObjectId()
        self.username = username
        self.password = password  # 密码加密后的字符串
        self.email = email
        self.phone = phone
        self.is_admin = is_admin
        self.avatar = avatar
        self.created_at = created_at or datetime.utcnow()
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "_id": self._id,
            "username": self.username,
            "password": self.password,
            "email": self.email,
            "phone": self.phone,
            "is_admin": self.is_admin,
            "avatar": self.avatar,
            "created_at": self.created_at
        }

class AISupplierModel:
    """AI供应商模型"""
    def __init__(self, name: str, api_key: str, api_url: str, description: Optional[str] = None,
                 is_active: bool = True, created_at: Optional[datetime] = None, _id: Optional[ObjectId] = None):
        self._id = _id or ObjectId()
        self.name = name  # 供应商名称，如kimi、deepseek等
        self.api_key = api_key  # API密钥
        self.api_url = api_url  # API请求地址
        self.description = description
        self.is_active = is_active
        self.created_at = created_at or datetime.utcnow()
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "_id": self._id,
            "name": self.name,
            "api_key": self.api_key,
            "api_url": self.api_url,
            "description": self.description,
            "is_active": self.is_active,
            "created_at": self.created_at
        }

class AIRoleModel:
    """AI角色模型"""
    def __init__(self, name: str, prompt: str, supplier_id: ObjectId, 
                 description: Optional[str] = None, avatar: Optional[str] = None,
                 is_active: bool = True, created_at: Optional[datetime] = None, 
                 model: Optional[str] = None, temperature: Optional[float] = None,
                 max_tokens: Optional[int] = None, role_status: Optional[str] = None,
                 is_default: Optional[bool] = None, _id: Optional[ObjectId] = None):
        self._id = _id or ObjectId()
        self.name = name  # 角色名称，如"儿科医生"、"心理咨询师"等
        self.prompt = prompt  # 角色提示词，定义AI的行为和知识范围
        self.supplier_id = supplier_id  # 关联的AI供应商ID
        self.description = description
        self.avatar = avatar  # 角色头像
        self.is_active = is_active
        self.created_at = created_at or datetime.utcnow()
        # 新增的AI参数
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.role_status = role_status
        self.is_default = is_default
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        result = {
            "_id": self._id,
            "name": self.name,
            "prompt": self.prompt,
            "supplier_id": self.supplier_id,
            "description": self.description,
            "avatar": self.avatar,
            "is_active": self.is_active,
            "created_at": self.created_at
        }
        
        # 添加可选的AI参数（只有在不为None时才添加）
        if self.model is not None:
            result["model"] = self.model
        if self.temperature is not None:
            result["temperature"] = self.temperature
        if self.max_tokens is not None:
            result["max_tokens"] = self.max_tokens
        if self.role_status is not None:
            result["status"] = self.role_status
        if self.is_default is not None:
            result["is_default"] = self.is_default
            
        return result

class ConsultationModel:
    """问诊会话模型（基本信息存储在MongoDB中，详细内容存储在Redis中）"""
    def __init__(self, user_id: ObjectId, title: str, role_id: ObjectId,
                 status: str = "active", created_at: Optional[datetime] = None, 
                 updated_at: Optional[datetime] = None, _id: Optional[ObjectId] = None):
        self._id = _id or ObjectId()
        self.user_id = user_id  # 用户ID
        self.title = title  # 会话标题
        self.role_id = role_id  # 使用的AI角色ID
        self.status = status  # active, closed
        self.created_at = created_at or datetime.utcnow()
        self.updated_at = updated_at or datetime.utcnow()
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "_id": self._id,
            "user_id": self.user_id,
            "title": self.title,
            "role_id": self.role_id,
            "status": self.status,
            "created_at": self.created_at,
            "updated_at": self.updated_at
        }

class MessageModel:
    """消息模型（存储在Redis中）"""
    def __init__(self, role: str, content: str, timestamp: Optional[datetime] = None):
        self.role = role  # user或assistant
        self.content = content
        self.timestamp = timestamp or datetime.utcnow()
    
    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式"""
        return {
            "role": self.role,
            "content": self.content,
            "timestamp": self.timestamp.isoformat() if self.timestamp else None
        }

# 集合名称定义
COLLECTION_NAMES = {
    "users": "users",
    "ai_suppliers": "ai_suppliers",
    "ai_roles": "ai_roles",
    "consultations": "consultations"
}